import gradio as gr
import numpy as np
from PIL import Image
from sklearn.cluster import KMeans
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from io import BytesIO


def extract_palette(image, n_colors=5):
    """
    使用 KMeans 聚类提取图片的主色调调色板

    参数:
        image: PIL Image 对象
        n_colors: 要提取的颜色数量，默认为 5

    返回:
        palette_image: 显示调色板的图片
        hex_colors_text: 十六进制颜色代码文本
    """
    if image is None:
        return None, "请先上传图像"

    # 将图像转换为 RGB 格式
    if image.mode != 'RGB':
        image = image.convert('RGB')

    # 将图像转换为 numpy 数组
    img_array = np.array(image)

    # 将图像重塑为像素列表 (height * width, 3)
    pixels = img_array.reshape(-1, 3)

    # 使用 KMeans 聚类找到主要颜色
    kmeans = KMeans(n_clusters=n_colors, random_state=42, n_init=10)
    kmeans.fit(pixels)

    # 获取聚类中心（主色调）
    colors = kmeans.cluster_centers_.astype(int)

    # 计算每个聚类的像素数量（用于排序）
    labels = kmeans.labels_
    label_counts = np.bincount(labels)

    # 按出现频率排序颜色
    sorted_indices = np.argsort(-label_counts)
    sorted_colors = colors[sorted_indices]

    # 创建调色板图像
    palette_image = create_palette_image(sorted_colors)

    # 生成十六进制颜色代码
    hex_colors = [rgb_to_hex(color) for color in sorted_colors]
    hex_colors_text = "\n".join(hex_colors)

    return palette_image, hex_colors_text


def rgb_to_hex(rgb):
    """将 RGB 颜色转换为十六进制格式"""
    return '#{:02x}{:02x}{:02x}'.format(int(rgb[0]), int(rgb[1]), int(rgb[2]))


def create_palette_image(colors):
    """创建调色板图像显示"""
    n_colors = len(colors)

    # 创建图形
    fig, ax = plt.subplots(figsize=(10, 2))

    # 绘制颜色条
    for i, color in enumerate(colors):
        rect = patches.Rectangle(
            (i, 0), 1, 1,
            facecolor=color / 255.0,
            edgecolor='none'
        )
        ax.add_patch(rect)

    # 设置坐标轴
    ax.set_xlim(0, n_colors)
    ax.set_ylim(0, 1)
    ax.axis('off')

    # 将图形转换为图像
    buf = BytesIO()
    plt.tight_layout(pad=0)
    plt.savefig(buf, format='png', bbox_inches='tight', dpi=150)
    buf.seek(0)
    palette_img = Image.open(buf)
    plt.close(fig)

    return palette_img


# 创建 Gradio 界面
with gr.Blocks(theme=gr.themes.Base(), css="""
    .gradio-container {
        background-color: #1a1f2e !important;
        color: #e0e0e0 !important;
    }
    body {
        background-color: #1a1f2e !important;
    }
    .main {
        background-color: #1a1f2e !important;
    }
    h1 {
        text-align: center;
        color: #ffffff !important;
        font-size: 24px !important;
        font-weight: 600 !important;
        margin-bottom: 10px !important;
    }
    .description {
        text-align: left;
        color: #e0e0e0 !important;
        font-size: 14px !important;
        margin-bottom: 20px !important;
    }
    .image-container, .output-container {
        background-color: #2a3244 !important;
        border-radius: 8px !important;
        padding: 15px !important;
    }
    label {
        color: #ffffff !important;
        font-size: 14px !important;
        font-weight: 500 !important;
    }
    .clear-button {
        background-color: #4a5568 !important;
        color: #ffffff !important;
        border: none !important;
        border-radius: 6px !important;
        font-weight: 500 !important;
    }
    .submit-button {
        background-color: #ff6b35 !important;
        color: #ffffff !important;
        border: none !important;
        border-radius: 6px !important;
        font-weight: 500 !important;
    }
    .flag-button {
        background-color: #6b7280 !important;
        color: #ffffff !important;
        border: none !important;
        border-radius: 6px !important;
        font-weight: 500 !important;
        width: 100% !important;
    }
    textarea {
        background-color: #1e2530 !important;
        color: #ffffff !important;
        border: 1px solid #3a4254 !important;
        border-radius: 6px !important;
        font-size: 16px !important;
        font-family: 'Courier New', Consolas, monospace !important;
        line-height: 1.6 !important;
        font-weight: 600 !important;
        letter-spacing: 0.5px !important;
    }
    .textbox_wrapper textarea {
        color: #ffffff !important;
    }
    input, textarea, .input-text {
        color: #ffffff !important;
    }
    .image-frame {
        background-color: #2a3244 !important;
        border: 1px solid #3a4254 !important;
        border-radius: 6px !important;
    }
    /* 强制所有文本输入框和文本区域使用白色文字 */
    .gr-textbox, .gr-text-input, .gr-textarea {
        color: #ffffff !important;
    }
    .gr-box textarea, .gr-box input {
        color: #ffffff !important;
        background-color: #1e2530 !important;
    }
    /* 确保占位符文本也可见 */
    textarea::placeholder, input::placeholder {
        color: #888888 !important;
    }
    /* 强制所有文本框内的值使用白色 */
    .gr-text-input textarea,
    .gr-textbox textarea,
    textarea[data-testid="textbox"],
    .gr-form textarea {
        color: #ffffff !important;
        -webkit-text-fill-color: #ffffff !important;
    }
    /* 针对特定的十六进制码文本框 */
    #hex-output-box textarea,
    #hex-output-box input,
    #hex-output-box {
        color: #ffffff !important;
        -webkit-text-fill-color: #ffffff !important;
    }
    /* 确保文本框内的文本值是白色的 */
    #hex-output-box .gr-text-input,
    #hex-output-box .gr-textbox {
        color: #ffffff !important;
    }
    /* 针对只读文本框 */
    textarea[readonly],
    textarea[disabled],
    .readonly textarea {
        color: #ffffff !important;
        -webkit-text-fill-color: #ffffff !important;
        opacity: 1 !important;
    }
""") as demo:

    gr.Markdown(
        "# KMeans调色板提取器"
    )
    
    gr.Markdown(
        "上传图像获取其调色板。",
        elem_classes="description"
    )

    with gr.Row():
        with gr.Column(scale=1):
            # 图像上传组件
            image_input = gr.Image(
                label="image",
                type="pil",
                height=300,
                elem_classes="image-container"
            )

            with gr.Row():
                clear_btn = gr.Button(
                    "Clear",
                    elem_classes="clear-button",
                    scale=1
                )
                submit_btn = gr.Button(
                    "Submit",
                    elem_classes="submit-button",
                    scale=1
                )

        with gr.Column(scale=1):
            # 调色板显示
            palette_output = gr.Image(
                label="",
                height=150,
                show_label=False,
                elem_classes="output-container"
            )

            # 十六进制颜色代码显示
            hex_output = gr.Textbox(
                label="色彩十六进制码",
                lines=5,
                interactive=False,
                elem_classes="output-container",
                elem_id="hex-output-box"
            )

            # Flag 按钮
            flag_btn = gr.Button(
                "Flag",
                elem_classes="flag-button"
            )

    # 绑定清除事件
    clear_btn.click(
        fn=lambda: (None, None, ""),
        inputs=[],
        outputs=[image_input, palette_output, hex_output]
    )
    
    # 绑定提交事件
    submit_btn.click(
        fn=extract_palette,
        inputs=[image_input],
        outputs=[palette_output, hex_output]
    )

    # 也可以在上传图像后自动处理
    image_input.change(
        fn=extract_palette,
        inputs=[image_input],
        outputs=[palette_output, hex_output]
    )


if __name__ == "__main__":
    demo.launch(share=False)
